import os
import json
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from skimage import measure
from pycocotools import mask as mask_util
import numpy as np
from segment_anything import SamPredictor, sam_model_registry, SamAutomaticMaskGenerator
from PIL import Image
import cv2
from typing import Optional, Union, List, Dict, Tuple


def get_sam_model(model_type: str = 'vit-h', checkpoint_path: str = '/data/yrz/pretrained/sam_vit_h_4b8939.pth', return_type='model', cuda=False):
    assert return_type in ['model', 'predictor', 'mask_generator']
    assert isinstance(model_type, str), f"model_type should be str, but is type{model_type}"
    assert isinstance(checkpoint_path, str), f"checkpoint_path should be str, but is type{checkpoint_path}"
    assert os.path.exists(checkpoint_path), f"please the path :{checkpoint_path} isn't exists"
    sam = sam_model_registry[model_type](checkpoint=checkpoint_path)
    if cuda: sam = sam.cuda()
    if return_type == 'model':
        return sam
    elif return_type == 'predictor':
        return SamPredictor(sam)
    elif return_type == 'mask_generator':
        return SamAutomaticMaskGenerator(sam, output_mode='coco_rle')
    else:
        raise ValueError(f'not support model_type:{model_type}')


def gen_auto_mask(img, sam=None, mask_generator=None):
    assert (sam is None) ^ (mask_generator is None), f'sam is {sam} and mask_generator is {mask_generator}'
    if mask_generator is None:
        mask_generator = SamAutomaticMaskGenerator(sam)
    masks = mask_generator.generate(img) # coco格式
    return masks


def gen_prompt_mask(img, sam=None, predictor=None, **kwargs):
    assert (sam is None) ^ (predictor is None), f'sam is {sam} and predictor is {predictor}'
    if predictor is None:
        predictor = SamPredictor(sam)
    predictor.set_image(img)
    masks, _, _ = predictor.predict(**kwargs)
    return masks


def open_img(img_path_: str) -> np.ndarray:
    img_ = cv2.imread(img_path_)
    if img_ is None:
        print(f"Could not load '{img_path_}' as an image, skipping...")
    img_ = cv2.cvtColor(img_, cv2.COLOR_BGR2RGB)

    return img_


# def visualize_img_(img_):
#     """
#     """
#     if len(img_.shape)==2:
#         plt.axis('off')
#         plt.imshow(img_.astype(np.uint8),cmap='gray',)
#     else:
#         plt.axis('off')
#         plt.imshow(img_.astype(np.uint8))
#
#     plt.show()

def normalize_array(
        arr: np.ndarray,
        target_range: Tuple[Union[int, float], Union[int, float]] = (0.0, 1.0),
        auto_dtype: bool = True,
        clip_outliers: bool = True,
) -> np.ndarray:
    """
    通用数组归一化函数，支持智能数据类型处理

    Args:
        arr: 输入数组 (支持任意维度)
        target_range: 目标数值范围，默认为 (0.0, 1.0)
        auto_dtype: 是否根据目标范围自动选择数据类型
            - 若target_range为整型范围且auto_dtype=True，自动转换为uint8
            - 若target_range包含浮点数，转换为float32
        clip_outliers: 是否将超出原始数据范围的值裁剪到目标区间

    Returns:
        归一化后的数组 (保留原始维度)

    Examples:
        # 浮点型数据归一化到 [0, 1]
        normalized = normalize_array(entropy_map)

        # 整型标签图归一化到 [0, 255]
        normalized = normalize_array(label_map, (0, 255))
    """
    # 输入校验
    if not isinstance(arr, np.ndarray):
        raise TypeError("Input must be a numpy array")
    if len(arr) == 0:
        raise ValueError("Input array is empty")

    # 处理全为常数的特殊情况
    if np.ptp(arr) == 0:
        return np.full_like(arr, target_range[0], dtype=np.float32)

    # 原始数据统计量
    src_min = np.nanmin(arr)
    src_max = np.nanmax(arr)

    # 线性归一化计算
    normalized = arr.astype(np.float32) - src_min
    normalized /= (src_max - src_min + 1e-8)  # 避免除以零

    # 缩放至目标区间
    scale = target_range[1] - target_range[0]
    normalized = normalized * scale + target_range[0]

    # 处理异常值
    if clip_outliers:
        normalized = np.clip(normalized, *sorted(target_range))

    # 自动数据类型转换逻辑
    if auto_dtype:
        if all(isinstance(x, int) for x in target_range):
            if target_range[1] - target_range[0] == 255:
                return normalized.astype(np.uint8)
            else:
                return normalized.astype(int)
        else:
            return normalized.astype(np.float32)

    return normalized


def visualize_3d_score(
        entropy_maps: Union[np.ndarray, List[np.ndarray]],
        title: Optional[Union[str, List[str]]] = None,
        figsize: tuple = (10, 6),
        elev: float = 30,
        azim: float = -120,
        cmap: str = 'viridis',
        stride: int = 1,
        axis_labels: tuple = ('Width', 'Height', 'Score'),
        colorbar: bool = True,
        surface_density: float = 1.0,
        save_path: Optional[str] = None,
        show: bool = True,
        **surface_kwargs
) -> None:
    """
    3D 熵值曲面可视化（输入为预计算熵值矩阵）

    Args:
        entropy_maps (Union[np.ndarray, List[np.ndarray]]): 熵值矩阵（单个或多个二维数组）
        title (Optional[Union[str, List[str]]]): 标题（与输入数量匹配）
        figsize (tuple): 画布尺寸，默认 (10,6)
        elev (float): 仰角，默认30度
        azim (float): 方位角，默认-120度
        cmap (str): 颜色映射，默认 'viridis'
        stride (int): 网格采样步长（控制密度），默认1（最高密度）
        axis_labels (tuple): 坐标轴标签，默认 ('Width','Height','Entropy')
        colorbar (bool): 是否显示颜色条，默认True
        surface_density (float): 曲面密度因子 (0.1~1.0)，通过插值提升密度，默认1.0（原始密度）
        save_path (Optional[str]): 保存路径
        show (bool): 是否显示图像，默认True
        **surface_kwargs: 额外曲面参数（如alpha=0.8设置透明度）
    """
    # 输入标准化
    if isinstance(entropy_maps, np.ndarray):
        entropy_maps = [entropy_maps]
    elif not isinstance(entropy_maps, list):
        raise TypeError("输入必须是 numpy 数组或数组列表")

    # 校验输入维度
    for i, emap in enumerate(entropy_maps):
        if emap.ndim != 2:
            raise ValueError(f"熵值矩阵 {i} 的维度应为 2D，实际为 {emap.shape}")
        if emap.dtype != np.float64 and emap.dtype != np.float32:
            raise TypeError(f"熵值矩阵 {i} 的数据类型应为 float，实际为 {emap.dtype}")

    # 创建画布
    fig = plt.figure(figsize=figsize)
    num_plots = len(entropy_maps)

    for idx, emap in enumerate(entropy_maps):
        # 处理曲面密度
        if surface_density != 1.0:
            from scipy.interpolate import griddata
            h, w = emap.shape
            x = np.linspace(0, w, int(w * surface_density))
            y = np.linspace(0, h, int(h * surface_density))
            xx, yy = np.meshgrid(x, y)
            points = np.array(np.meshgrid(np.arange(w), np.arange(h))).T.reshape(-1, 2)
            emap = griddata(points, emap.flatten(), (xx, yy), method='cubic')

        # 生成网格
        h, w = emap.shape
        x = np.arange(w)
        y = np.arange(h)
        xx, yy = np.meshgrid(x, y)

        # 创建3D子图
        ax = fig.add_subplot(1, num_plots, idx + 1, projection='3d')
        surf = ax.plot_surface(
            xx, yy, emap,
            cmap=cmap,
            rstride=stride,
            cstride=stride,
            antialiased=False,
            **surface_kwargs
        )

        # 视图设置
        ax.view_init(elev=elev, azim=azim)
        ax.set_xlabel(axis_labels[0], labelpad=10)
        ax.set_ylabel(axis_labels[1], labelpad=10)
        ax.set_zlabel(axis_labels[2], labelpad=10)

        # 标题与颜色条
        if title:
            if isinstance(title, list) and len(title) == num_plots:
                ax.set_title(title[idx])
            else:
                ax.set_title(title if num_plots == 1 else f"{title} ({idx + 1})")
        if colorbar:
            fig.colorbar(surf, ax=ax, shrink=0.6, aspect=10)

    # 保存与显示
    if save_path:
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        plt.savefig(save_path, bbox_inches='tight', dpi=300)
    if show:
        plt.show()
    plt.close()


def visualize_img_(
        img_: Union[np.ndarray, List[np.ndarray]],
        title: Optional[str] = None,
        cmap: Optional[str] = None,
        palette: Optional[Union[List, np.ndarray, Dict]] = None,
        save_path: Optional[str] = None,
        figsize: Optional[tuple] = None,
        axis_off: bool = True,
        show: bool = False,
        **kwargs
) -> None:
    """
    Visualize one or multiple images using matplotlib.

    Args:
        img_ (Union[np.ndarray, List[np.ndarray]]): Input image(s) as a numpy array or a list of numpy arrays.
        title (Optional[str]): Title of the image(s). Default: None.
        cmap (Optional[str]): Color map for single-channel images. Default: None (auto-detected).
        palette (Optional[Union[List, np.ndarray, Dict]]): Color palette for single-channel images.
            If provided, it will override `cmap`. Can be:
            - A list of RGB colors (e.g., [[255, 0, 0], [0, 255, 0], ...]).
            - A numpy array of shape (N, 3) where N is the number of colors.
            - A dictionary mapping labels to RGB colors (e.g., {0: [255, 0, 0], 1: [0, 255, 0], ...}).
            Default: None.
        save_path (Optional[str]): Path to save the image. If None, the image will not be saved. Default: None.
        figsize (Optional[tuple]): Figure size (width, height) in inches. Default: None (auto-detected).
        axis_off (bool): Whether to hide the axes. Default: True.
        show (bool): Whether to display the image(s). Default: True.
        **kwargs: Additional arguments to pass to `plt.imshow`.

    Returns:
        None
    """
    # Ensure input is a list for uniform processing
    if isinstance(img_, np.ndarray):
        img_ = [img_]
    elif not isinstance(img_, list):
        raise TypeError("Input must be a numpy array or a list of numpy arrays.")

    # Determine the number of images
    num_images = len(img_)
    if num_images == 0:
        raise ValueError("Input list is empty.")

    # Set up the figure and axes
    if figsize is None:
        figsize = (6 * num_images, 6)  # Default size for multiple images
    plt.figure(figsize=figsize)
    res = []
    for i, img in enumerate(img_):
        if not isinstance(img, np.ndarray):
            raise TypeError(f"Element {i} in the list is not a numpy array.")

        # Handle single-channel images
        if len(img.shape) == 2:
            if palette is not None:
                # Convert single-channel image to pseudo-color using palette
                img = apply_palette(img, palette)
                cmap = None  # Disable cmap since we are using palette
            elif cmap is None:
                cmap = 'gray'  # Default colormap for grayscale images
        elif len(img.shape) != 3:
            raise ValueError(f"Image {i} has invalid shape: {img.shape}. Expected 2D or 3D array.")

        # Create subplot for each image
        plt.subplot(1, num_images, i + 1)
        plt.imshow(img, cmap=cmap, **kwargs)
        if title:
            plt.title(title if num_images == 1 else f"{title} ({i + 1})")
        if axis_off:
            plt.axis('off')
        res.append(img)

    # Save the image if save_path is provided
    if save_path:
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        plt.savefig(save_path, bbox_inches='tight', pad_inches=0)
        print(f"Image saved to: {save_path}")

    # Display the image(s)
    if show:
        plt.show()

    # Close the figure to free memory
    plt.close()
    return res


def apply_palette(img: np.ndarray, palette: Union[List, np.ndarray, Dict]) -> np.ndarray:
    """
    Apply a color palette to a single-channel image.

    Args:
        img (np.ndarray): Single-channel image of shape (H, W).
        palette (Union[List, np.ndarray, Dict]]): Color palette.

    Returns:
        np.ndarray: Pseudo-color image of shape (H, W, 3).
    """
    if isinstance(palette, dict):
        # Convert dictionary palette to a list of colors
        max_label = max(palette.keys())
        palette_list = np.zeros((max_label + 1, 3), dtype=np.uint8)
        for label, color in palette.items():
            palette_list[label] = color
        palette = palette_list
    elif isinstance(palette, list):
        palette = np.array(palette, dtype=np.uint8)

    if not isinstance(palette, np.ndarray):
        raise TypeError("Palette must be a list, numpy array, or dictionary.")

    # Ensure palette is of shape (N, 3)
    if palette.shape[1] != 3:
        raise ValueError("Palette must have 3 color channels (RGB).")

    # Create pseudo-color image
    h, w = img.shape
    color_img = np.zeros((h, w, 3), dtype=np.uint8)
    for label, color in enumerate(palette):
        color_img[img == label] = color

    return color_img


def random_color():
    return np.random.rand(3, )


def gen_comb_mask(masks):
    overlay = np.zeros(masks[0]['segmentation']['size'])
    for mask_data in masks:
        mask = mask_util.decode(mask_data['segmentation'])  # 布尔数组
        overlay[mask == 1] = 1
    return overlay


class RandomMaskSampler:
    """
    基于累加值控制的随机采样器，支持 COCO 格式的 masks 数据，使用 NumPy 实现。
    """

    def __init__(self, masks, weight_key="stability_score"):
        """
        初始化采样器。

        参数:
            masks (list): 包含 COCO 格式的 masks 数据。
            weight_key (str): 用于加权采样的字段，例如 "stability_score"。
        """
        self.masks = masks
        self.weight_key = weight_key

        # 提取面积和权重
        self.areas = np.array([mask["area"] for mask in masks])
        self.weights = np.array([mask.get(weight_key, 1.0) for mask in masks])

        # 计算总图像面积
        image_width, image_height = masks[0]["segmentation"]["size"]
        self.total_image_area = image_width * image_height

        # 预计算权重总和，避免重复计算
        self.weight_sum = self.weights.sum()

    def sample(self, target_area_ratio):
        """
        执行一次随机采样，按累加值控制。

        参数:
            target_area_ratio (float): 面积目标比例 (0~1)。

        返回:
            list: 选中的 mask 数据。
        """
        target_area = target_area_ratio * self.total_image_area
        cumulative_area = 0

        # 剩余可选的索引和布尔掩码
        mask = np.ones(len(self.masks), dtype=bool)
        selected_indices = []

        while cumulative_area < target_area and mask.any():
            # 按权重计算概率
            probabilities = self.weights[mask] / self.weights[mask].sum()

            # 随机采样
            chosen_index = np.random.choice(np.where(mask)[0], p=probabilities)

            # 更新选择
            selected_indices.append(chosen_index)
            cumulative_area += self.areas[chosen_index]

            # 移除已选索引
            mask[chosen_index] = False
        # probabilities = self.weights / self.weights.sum()
        # selected_indices = np.random.choice(int(len(self.masks)), int(len(self.masks) / 2), replace=False,p=probabilities)

        return [self.masks[int(i)] for i in selected_indices]

    def generate_combinations(self, num_combinations=10):
        """
        生成多个采样组合。

        参数:
            num_combinations (int): 生成的随机组合数量。

        返回:
            list: 包含多个采样组合的列表，每个组合是一个 mask 列表。
        """
        return [self.sample(target_area_ratio=np.random.random() * 0.2 + 0.4) for _ in range(num_combinations)]


def visualize_masks(masks, image=None, ):
    if image is not None:
        overlay = image.copy()
    else:
        overlay = np.ones((*masks[0]['segmentation']['size'], 3)) * 255

    fig, ax = plt.subplots(figsize=(12, 12))

    # 为每个mask分配一个颜色，并显示
    for mask_data in masks:
        mask = mask_util.decode(mask_data['segmentation'])  # 布尔数组

        # 为每个实例分配一个颜色
        color = random_color()

        # 为当前 mask 区域着色
        for c in range(3):  # RGB 通道
            overlay[..., c] = np.where(mask, color[c] * 255, overlay[..., c])

        # 为每个mask生成黑色边界
        contours = measure.find_contours(mask, 0.5)
        for contour in contours:
            plt.plot(contour[:, 1], contour[:, 0], linewidth=0.5, color='black')  # 黑色边界线

    ax.imshow(overlay.astype(np.uint8), alpha=0.5)  # 确保数据为整数类型以显示
    ax.axis('off')  # 关闭坐标轴
    plt.show()


if __name__ == '__main__':
    from tqdm import tqdm

    # img_path = f'/data/yrz/repos/ST-DASegNet/demo/demo.png'
    # img = open_img(img_path)
    # visualize_img_(img)
    # model_type = 'vit_h'  # ["<model_type>"]
    # checkpoint_path = '/data/yrz/pretrained/sam_vit_h_4b8939.pth'  # ["<path/to/checkpoint>"]
    # mask_generator = get_sam_model(model_type=model_type, checkpoint_path=checkpoint_path, return_type='mask_generator',cuda=True)
    #
    # masks = gen_auto_mask(img, mask_generator=mask_generator)
    # visualize_masks(masks,image=img)
    # print(masks)

    json_path = f'/data/yrz/repos/ST-DASegNet/data/Vaihingen_IRRG_DA/auto_mask_dir/train/area11_512_1024_1024_1536.json'
    with open(json_path, 'r') as f:
        masks = json.load(f)
    visualize_masks(masks)

    comb_mask = gen_comb_mask(masks) * 255
    visualize_img_(comb_mask)

    mask_sampler = RandomMaskSampler(masks)
    select_masks = mask_sampler.sample(0.5)
    visualize_masks(select_masks)
    comb_mask = gen_comb_mask(select_masks) * 255
    visualize_img_(comb_mask)

    combinations = mask_sampler.generate_combinations(num_combinations=50)
    for i, comb in enumerate(tqdm(combinations)):
        visualize_img_(gen_comb_mask(comb))
        print(f'process {i}')
